package com.tony.jdbc.dynamic.service.impl;

import com.tony.jdbc.dynamic.datasource.DynamicDataSourceKeyHolder;
import com.tony.jdbc.dynamic.service.IDynamicService;
import com.tony.jdbc.dynamic.service.ITransactionService;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.ApplicationContext;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.jdbc.core.namedparam.NamedParameterJdbcTemplate;
import org.springframework.stereotype.Service;

import java.util.function.Consumer;
import java.util.function.Function;

/**
 * @author tony
 * @desc 动态数据源服务接口实现类，默认使用spring事务。
 * 1.先调用DynamicDataSourceKeyHolder.setKey(datasourceName)方法切换数据源。
 * 2.调用transactionService.execute方法执行事务。
 * 3.最后移除数据源。
 * @date 2025/11/18 21:00
 */
@Service
public class DynamicServiceImpl implements IDynamicService {
    @Value("${spring.jta.enabled:false}")
    private boolean jta;
    private final ApplicationContext applicationContext;
    /**
     * 事务服务，默认使用spring事务
     */
    private final ITransactionService transactionService;
    /**
     * 数据库操作对象
     */
    private final JdbcTemplate jdbc;
    /**
     * 支持绑定参数的数据库操作对象
     */
    private final NamedParameterJdbcTemplate namedParameterJdbc;

    public DynamicServiceImpl(ApplicationContext applicationContext, ITransactionService transactionService, JdbcTemplate jdbc, NamedParameterJdbcTemplate namedParameterJdbc) {
        this.applicationContext = applicationContext;
        this.transactionService = transactionService;
        this.jdbc = jdbc;
        this.namedParameterJdbc = namedParameterJdbc;
    }

    /**
     * 获取当前数据源的JdbcTemplate对象
     *
     * @param datasourceName 数据源名称
     * @return JdbcTemplate
     */
    private JdbcTemplate getCurrentJdbcTemplate(String datasourceName) {
        JdbcTemplate currentJdbc;
        if (jta) {
            currentJdbc = applicationContext.getBean(datasourceName + DynamicDataSourceKeyHolder.JDBC_TEMPLATE, JdbcTemplate.class);
        } else {
            currentJdbc = this.jdbc;
        }
        if (currentJdbc == null) {
            throw new RuntimeException("JdbcTemplate is null");
        }
        return currentJdbc;
    }

    /**
     * 获取当前数据源的NamedParameterJdbcTemplate对象
     *
     * @param datasourceName 数据源名称
     * @return JdbcTemplate
     */
    private NamedParameterJdbcTemplate getCurrentNamedParameterJdbcTemplate(String datasourceName) {
        NamedParameterJdbcTemplate currentJdbc;
        if (jta) {
            currentJdbc = applicationContext.getBean(datasourceName + DynamicDataSourceKeyHolder.NAMED_JDBC_TEMPLATE, NamedParameterJdbcTemplate.class);
        } else {
            currentJdbc = this.namedParameterJdbc;
        }
        if (currentJdbc == null) {
            throw new RuntimeException("NamedParameterJdbcTemplate is null");
        }
        return currentJdbc;
    }

    @Override
    public void executeJdbc(String datasourceName, Consumer<JdbcTemplate> consumer) {
        //获取当前数据源标识
        String originKey = DynamicDataSourceKeyHolder.getKey();
        DynamicDataSourceKeyHolder.setKey(datasourceName);
        JdbcTemplate currentJdbc = getCurrentJdbcTemplate(datasourceName);
        try {
            transactionService.execute(() -> consumer.accept(currentJdbc));
        } finally {
            //DynamicDataSourceKeyHolder.removeKey();
            //恢复数据源标识，保证执行中可进行多个数据源切换
            DynamicDataSourceKeyHolder.setKey(originKey);
        }
    }


    @Override
    public <T> T executeJdbcWithResult(String datasourceName, Function<JdbcTemplate, T> consumer) {
        //获取当前数据源标识
        String originKey = DynamicDataSourceKeyHolder.getKey();
        DynamicDataSourceKeyHolder.setKey(datasourceName);
        JdbcTemplate currentJdbc = getCurrentJdbcTemplate(datasourceName);
        try {
            return transactionService.execute(() -> consumer.apply(currentJdbc));
        } finally {
            //DynamicDataSourceKeyHolder.removeKey();
            //恢复数据源标识，保证执行中可进行多个数据源切换
            DynamicDataSourceKeyHolder.setKey(originKey);
        }
    }

    @Override
    public void executeParamJdbc(String datasourceName, Consumer<NamedParameterJdbcTemplate> consumer) {
        //获取当前数据源标识
        String originKey = DynamicDataSourceKeyHolder.getKey();
        DynamicDataSourceKeyHolder.setKey(datasourceName);
        NamedParameterJdbcTemplate currentJdbc = getCurrentNamedParameterJdbcTemplate(datasourceName);
        try {
            transactionService.execute(() -> consumer.accept(currentJdbc));
        } finally {
            //DynamicDataSourceKeyHolder.removeKey();
            //恢复数据源标识，保证执行中可进行多个数据源切换
            DynamicDataSourceKeyHolder.setKey(originKey);
        }
    }

    @Override
    public <T> T executeParamJdbcWithResult(String datasourceName, Function<NamedParameterJdbcTemplate, T> function) {
        //获取当前数据源标识
        String originKey = DynamicDataSourceKeyHolder.getKey();
        DynamicDataSourceKeyHolder.setKey(datasourceName);
        NamedParameterJdbcTemplate currentJdbc = getCurrentNamedParameterJdbcTemplate(datasourceName);
        try {
            return transactionService.execute(() -> function.apply(currentJdbc));
        } finally {
            //DynamicDataSourceKeyHolder.removeKey();
            //恢复数据源标识，保证执行中可进行多个数据源切换
            DynamicDataSourceKeyHolder.setKey(originKey);
        }
    }

    @Override
    public int updateJdbc(String datasourceName, Function<JdbcTemplate, Integer> function) {
        //获取当前数据源标识
        String originKey = DynamicDataSourceKeyHolder.getKey();
        DynamicDataSourceKeyHolder.setKey(datasourceName);
        JdbcTemplate currentJdbc = getCurrentJdbcTemplate(datasourceName);
        try {
            return transactionService.execute(() -> function.apply(currentJdbc));
        } finally {
            //DynamicDataSourceKeyHolder.removeKey();
            //恢复数据源标识，保证执行中可进行多个数据源切换
            DynamicDataSourceKeyHolder.setKey(originKey);
        }
    }

    @Override
    public int updateParamJdbc(String datasourceName, Function<NamedParameterJdbcTemplate, Integer> function) {
        //获取当前数据源标识
        String originKey = DynamicDataSourceKeyHolder.getKey();
        DynamicDataSourceKeyHolder.setKey(datasourceName);
        NamedParameterJdbcTemplate currentJdbc = getCurrentNamedParameterJdbcTemplate(datasourceName);
        try {
            return transactionService.execute(() -> function.apply(currentJdbc));
        } finally {
            //DynamicDataSourceKeyHolder.removeKey();
            //恢复数据源标识，保证执行中可进行多个数据源切换
            DynamicDataSourceKeyHolder.setKey(originKey);
        }
    }
}
